from flask import Flask, render_template, request, send_from_directory
import os
from werkzeug.utils import secure_filename

app = Flask(__name__)

# 配置上传文件夹和允许的文件扩展名
UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

# 检查文件扩展名是否允许
def allowed_file(filename):
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

# 首页路由
@app.route('/')
def upload_form():
    return render_template('upload.html')

# 上传文件处理路由
@app.route('/', methods=['POST'])
def upload_image():
    if 'file' not in request.files:
        return render_template('upload.html', message='No file part')
    file = request.files['file']
    if file.filename == '':
        return render_template('upload.html', message='No image selected for uploading')
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(file_path)

        # 在这里插入你的图像处理逻辑，现在只是简单地复制文件
        # 比如，mask_filename = your_image_processing_function(file_path)
        mask_filename = filename  # 暂时用原图代替mask图

        return render_template('upload.html', filename=filename, mask_filename=mask_filename)
    else:
        return render_template('upload.html', message='Allowed image types are -> png, jpg, jpeg')

# 服务图片文件
@app.route('/display/<filename>')
def display_image(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)

if __name__ == "__main__":
    app.run()
